from __future__ import division
import numpy as np
import scipy.constants as sconst
import scipy.linalg
import pickle

global sm

class SpinMatrix():
    def __init__(self):
        self.sx = {}
        self.sy = {}
        self.sz = {}
    
    def Sx(self, dim):
        if str(dim) in self.sx.keys():
            return self.sx[str(dim)]
        else:
            sx = np.diag(np.sqrt(dim * (dim+1) - np.arange(-dim, dim) * (np.arange(-dim+1, dim+1)))/2, 1)
            sx += sx.T.conj()
            self.sx[str(dim)] = sx
            return sx
            
    def Sy(self, dim):
        if str(dim) in self.sy.keys():
            return self.sy[str(dim)]
        else:
            sy = np.diag(-1j * np.sqrt(dim * (dim+1) - np.arange(-dim, dim) * (np.arange(-dim+1, dim+1)))/2, 1)
            sy += sy.T.conj()
            self.sy[str(dim)] = sy
            return sy
            
    def Sz(self, dim):
        if str(dim) in self.sz.keys():
            return self.sz[str(dim)]
        else:
            sz = np.diag(- np.arange(-dim, dim+1))
            self.sz[str(dim)] = sz
            return sz

class System():
    def __init__(self, B=0):
        self.B=B
        self.SpinList = []
        self.get_dim_list()
        self.n_spins = 0
        self.zfs=True
        self.zee=True
        self.dipolar=True
        self.indexlib = {'zeeman' : [], 'zfs' : [], 'dipolar' : []}
        self.create_IndexList()
        self.hamiltonian()
        self.rho_init()
        
    def rho_init(self):
        self.rho=1
        for spin in self.SpinList:
            self.rho=np.kron(self.rho, spin.state)
        
    def get_dim_list(self):
        self.dim=[]
        for spin in self.SpinList:
            self.dim.append(2*spin.art.spinNumber+1)
        self.dim_tot=np.prod(np.array(self.dim))
    
    def add_spin(self, spin):
        spin.system = True
        self.SpinList.append(spin)
        self.n_spins += 1
        self.get_dim_list()
        self.indexlib['zeeman'] += [self.n_spins-1]
        self.indexlib['zfs'] += [self.n_spins-1]
        self.indexlib['dipolar'] += [[i, self.n_spins-1] for i in range(self.n_spins-1)]
        self.hamiltonian()
        #self.SingleSpinIndexList()

    def delete_spin(self, spin):
        spin_index = self.SpinList.index(spin)
        spin.system = False
        self.SpinList.remove(spin)
        self.n_spins -= 1
        self.get_dim_list()
        if spin_index in self.indexlib['zeeman']:
            self.indexlib['zeeman'].remove(spin_index)
        if spin_index in self.indexlib['zfs']:
            self.indexlib['zfs'].remove(spin_index)
        for pair in self.indexlib['dipolar']:
            if spin_index in pair:
                self.indexlib['dipolar'].remove(pair)
        for itype in self.indexlib.values():
            for i in range(len(itype)):
                try:
                    for j in range(len(itype[i])):
                        if itype[i][j] > spin_index:
                            itype[i][j] -= 1
                except:
                    if itype[i] > spin_index:
                        itype[i] -= 1
        #self.SingleSpinIndexList()
    
    """def hamiltonian(self):
        if self.n_spins>0:
            self.H=H_zee.H(self, self.SingleSpinIndexList)+H_zfs.H(self, self.SingleSpinIndexList)
            if self.n_spins>1:
                self.H+=H_dip(self, self.AllInteractSpinIndex)
            self.energy, self.eigenstate=np.linalg.eigh(self.H)
        else:
            self.H=np.array([])
            self.energy, self.eigenstate=[], []"""
        
    def hamiltonian(self):
        self.H = np.zeros((self.dim_tot, self.dim_tot), dtype = complex)
        for name, i in self.indexlib.items():
            if len(i) > 0:
                self.H += H_lib[name].H(self, i)
        self.energy, self.eigenstate=np.linalg.eigh(self.H)

    def Sx(self, i):
        dim_before=np.prod(self.dim[:i])
        dim_after=np.prod(self.dim[i+1:])
        sx=np.kron(np.kron(np.eye(dim_before), sm.Sx(self.SpinList[i].art.spinNumber)), np.eye(dim_after))
        return sx

    def Sy(self, i):
        dim_before=np.prod(self.dim[:i])
        dim_after=np.prod(self.dim[i+1:])
        sy=np.kron(np.kron(np.eye(dim_before), sm.Sy(self.SpinList[i].art.spinNumber)), np.eye(dim_after))
        return sy

    def Sz(self, i):
        dim_before=np.prod(self.dim[:i])
        dim_after=np.prod(self.dim[i+1:])
        sz=np.kron(np.kron(np.eye(dim_before), sm.Sz(self.SpinList[i].art.spinNumber)), np.eye(dim_after))
        return sz    

    def U_0(self, t):
        u=scipy.linalg.expm2(-1j*2*np.pi*self.H*t)
        u_i=scipy.linalg.expm2(1j*2*np.pi*self.H*t)
        return u, u_i

    def write_vec(self, vec):
        ind = np.nonzero(np.abs(vec) >= 0.001)[0]
        text = ''
        for i in ind:
            state = ''
            dimdel = 0
            for j in range(len(self.dim)):
                state += str(self.SpinList[j].art.spinNumber-int(((i)-dimdel)/np.prod(self.dim[j+1:]))) + ','
                dimdel += int(np.prod(self.dim[j+1:])*int((i-dimdel)/np.prod(self.dim[j+1:])))
            text += str(complex(round(vec[i].real,3),round(vec[i].imag,3))) + '|' + state[:-1] + '>' + '+'
        return text[:-1]
        
    """def write_vec(self, vec):
        l1=range(len(vec))
        n=self.n_spins
        matrix=vec.reshape((self.dim-1)/2)
        text=""
        for i in list(matrix):
            if i!=0:
                pos=i.index()
                for k in range(n):
                    state=pos[k]-self.SpinList[k].spinNumber
                    text+=str(state)+", "
        return text"""

    def create_IndexList(self):
        self.SingleSpinIndexList=range(len(self.SpinList))
        self.AllInteractSpinIndex=[]
        for i1 in range(len(self.SpinList)-1):
            for i2 in range(i1+1, len(self.SpinList)):
                self.AllInteractSpinIndex.append([i1, i2])
            
                
class pulse():
    def __init__(self, system, state1, state2, alpha=0, phi=0, omega_t=0):
        self.state1=state1
        self.state2=state2
        self.rotation_matrix()
        self.alpha=alpha
        self.phi=phi
        self.phi_real=phi+omega_t
        self.U()

    def rotation_matrix(self):
        self.S_rot=0.5*(np.exp(-1j*self.phi_real)*np.outer(self.state1, self.state2)+np.exp(1j*self.phi_real)*np.outer(self.state2, self.state1))

    def U(self):
        self.U=scipy.linalg.expm2(-1j*2*np.pi*alpha*self.S_rot)
        self.U_i=scipy.linalg.expm2(1j*2*np.pi*alpha*self.S_rot)

    def change_state(self, new_state1, new_state2):
        self.state1=new_state1
        self.state2=new_state2
        self.rotation_matrix()
        self.U()
        
        
class wait():
    def __init__(self, system, tau):
        self.tau=tau
        self.H=system.H
        self.U()

    def U(self):
        self.U=scipy.linalg.expm2(-1j*2*np.pi*self.H*tau)
        self.U_i=scipy.linalg.expm2(1j*2*np.pi*self.H*tau)
        

class sequence():
    def __init__(self, tau=1):
        self.UList=[]
        self.calculate_U()
        self.tau=tau

    def add_part(self, part):
        self.UList.append(part)
        self.U=np.dot(part.U, self.U)
        self.U_i=np.dot(self.U_i, part.U_i)
        
    def calculate_U(self):
        for part in UList:
            self.U=np.dot(part.U, self.U)
            self.U_i=np.dot(self.U_i, part.U_i)

    def delete_part(self, i):
        del(self.UList[i])
        self.calculate_U()

    def save_sequence(self, filename="sequence0"):
        data="***\t "+filename+" \t ***\n"
        for part in UList:
            data+="----------------------------------------------------------\n"
            if part.__class__.__name__=="pulse":
                data+="Pulse\t"+str(part.alpha/np.pi)+"\t"+str(part.phi)+"\n"
            if part.__class__.__name__=="wait":
                data+="Wait\t"+str(part.tau/self.tau)+"\n"
        data+="***********************************************************\n"
        f=open(filename, "w")
        f.write(data)
        f.close()

    def load_sequence(self, filename, system):
        f=open(filename, "r")
        data=f.readlines()
        f.close()
        for line in data:
            if line[0]=="Pulse":
                alpha=float(line[1])
                phi=float(line[2])
                p=pulse(system, system.vec[0], system.vec[1], alpha, phi)
                self.add_part(p)
            elif line[0]=="Wait":
                tau=float(line[1])*self.tau
                w=wait(system, tau)
                self.add_part(w)

    def change_tau(self, new_tau):
        for part in self.UList:
            if part.__class__.__name__=="wait":
                part.tau=part.tau/self.tau*new_tau
                self.tau=new_tau
                part.U()
        self.calculate_U()
        
    def change_position(self, pos_old, pos_new):
        part=self.UList[pos_old]
        del(self.UList[pos_old])
        self.UList.insert(part, pos_new)
        self.calculate_U()

sm=SpinMatrix()

class DynamicalDecoupling(sequence):
    def __init__(self, tau=1):
        sequence.__init__(self, tau)
        self.n=0
        self.UListStart=[]
        self.UListEnd=[]
        self.UListMiddle=[]
        self.calculate_t()
        philist=[]
        self.t_function="tau"
        phi_period=[0]

    def calculate_U(self):
        self.UList.extend(UList.Start)
        for i in range(n):
            pass
            
    def calculate_t(self):
        self.t_function.replace("tau", "self.tau")
        try:
            tlist=[eval(self.t_function) for t in range(n)]
        except:
            print "not a correct function for time"

    def calculate_phi(self):
        pass
            
        

def rabi(system, state_number1, state_number2, f_rabi):
    f_f=system.val[state_number1]-system.val[state_number2]
    pass


class Hamiltonian():
    def __init__(self):
        self.name="H"
        self.equation=[]
        self.SpinNumber=1

    def H(self, system, SpinIndex, B=np.zeros(3)):
        #change self.H or def H to another name or think about saving the Hamiltonian-Matrixes.
        H=np.zeros((system.dim_tot, system.dim_tot), dtype = complex)
        if self.SpinNumber==1:
            for index in SpinIndex:
                Sx=system.Sx(index)
                Sy=system.Sy(index)
                Sz=system.Sz(index)
                spin=system.SpinList[index]
                for term in self.equation:
                    term_H=np.eye(system.dim_tot)
                    for factor in term:
                        if "^" in factor:
                            power=eval(factor.split("^")[1])
                            factor=factor.split("^")[0]
                        else:
                            power=1
                        if Dict_H.has_key(factor):
                            factor=eval(Dict_H[factor])
                        else: factor=eval(factor)
                        if np.array(factor).size==1:
                            factor=factor**power*np.eye(system.dim_tot)
                        term_H=np.dot(term_H, factor)
                    H+=term_H
        elif self.SpinNumber==2:
            for index1, index2 in SpinIndex:
                spin1=system.SpinList[index1]
                Sx1=system.Sx(index1)
                Sy1=system.Sy(index1)
                Sz1=system.Sz(index1)
                spin2=system.SpinList[index2]
                Sx2=system.Sx(index2)
                Sy2=system.Sy(index2)
                Sz2=system.Sz(index2)
                for term in self.equation:
                    term_H=np.eye(system.dim_tot)
                    for factor in term:
                        if "^" in factor:
                            power=eval(factor.split("^")[1])
                            factor=factor.split("^")[0]
                        else:
                            power=1
                        if Dict_H.has_key(factor):
                            factor=eval(Dict_H[factor])
                        else: factor=eval(factor)
                        if np.array(factor).size==1:
                            factor=factor**power*np.eye(system.dim_tot)
                        term_H=np.dot(term_H, factor)
                    H+=term_H
        return H

    def save_Hamiltonian(self, filename):
        try:
            f=open(filename, "r")
            text=f.read()
            f.close()
        except:
            text=""
        text+="+++++++++++++++++++++++++++++++++++++++++++\n"
        text+="***\t"+self.name+"\t( "+str(self.SpinNumber)+" Spins )"+"\t***\n"
        text+="-------------------------------------------\n"
        for term in self.equation[:-1]:
            for factor in term[:-1]:
                text+=factor
                text+="*"
            text+=term[-1]
            text+="+"
        for factor in self.equation[-1][:-1]:
            text+=factor
            text+="*"
        text+=self.equation[-1][-1]
        text+="\n ___________________________________________\n"
        f=open(filename, "w")
        f.write(text)
        f.close()
            
        
    def load_Hamiltonian(self, system, text):
        attributes=text[1].split("\t")
        self.name=attributes[1]
        self.SpinNumber=attributes[2].split(" ")[1]
        equ=text[3]
        equ=equ.replace("+", "\t")
        equ=equ.replace("*", " ")
        equ=equ=equ.split("\t")
        for i in range(len(equ)):
            equ[i]=equ[i].split(" ")
        self.equation=equ

        
    def load_Hamiltonian_from_dic(self, dic, name, system):
        load_Hamiltonian(system, dic[name])
    
    """
    def all_interact(self, system):
        self.SpinIndex=[]
        for i1 in range(len(system.SpinList)-1):
            for i2 in range(i1+1, len(system.SpinList)):
                self.SpinIndex.append([i1, i2])
    """
                        
Dict_H={"gyro":"spin.art.gyro", "gyro1":"spin1.art.gyro", "gyro2":"spin2.art.gyro",
        "Bx":"B[0]", "By":"B[1]", "Bz":"B[2]", "Sx":"Sx", "Sy":"Sy", "Sz":"Sz",
        "D_zfs":"spin.art.zfs", "D_zfs1":"spin1.art.zfs", "D_zfs2":"spin2.art.zfs",
        "r12":"spin1.get_distance(spin2)", "S":"spin.art.spinNumber", "h":"sconst.h",
        "S1":"spin1.art.spinNumber", "S2":"spin2.art.spinNumber", "mu0":"sconst.mu_0",
        "orientation_x":"spin.orientation_axis[0]", "orientation_y":"spin.orientation_axis[1]",
        "orientation_z":"spin.orientation_axis[2]", "orientation_x1":"spin1.orientation_axis[0]",
        "orientation_y1":"spin1.orientation_axis[1]", "orientation_z1":"spin1.orientation_axis[2]",
        "orientation_x2":"spin2.orientation_axis[0]", "orientation_y2":"spin2.orientation_axis[1]",
        "orientation_z2":"spin2.orientation_axis[2]",
        "relative_x":"spin1.get_relative_coordinates(spin2)[0]",
        "relative_y":"spin1.get_relative_coordinates(spin2)[1]",
        "relative_z":"spin1.get_relative_coordinates(spin2)[2]"}
        
Dict_H1 = {"gyro":"spin.art.gyro", "Bx":"B[0]", "By":"B[1]", "Bz":"B[2]", 
           "Sx":"Sx", "Sy":"Sy", "Sz":"Sz", "D_zfs":"spin.art.zfs", 
           "S":"spin.art.spinNumber", "h":"sconst.h", "mu0":"sconst.mu_0", 
           "orientation_x":"spin.orientation_axis[0]", 
           "orientation_y":"spin.orientation_axis[1]", 
           "orientation_z":"spin.orientation_axis[2]"}
           
Dict_H2 = {"gyro1":"spin1.art.gyro", "gyro2":"spin2.art.gyro", 
           "Bx":"B[0]", "By":"B[1]", "Bz":"B[2]", "Sx1":"Sx1", "Sy1":"Sy1", "Sz1":"Sz1", 
           "Sx2":"Sx2", "Sy2":"Sy2", "Sz2":"Sz2", "D_zfs1":"spin1.art.zfs", "D_zfs2":"spin2.art.zfs",
           "r12":"spin1.get_distance(spin2)", "h":"sconst.h", "S1":"spin1.art.spinNumber",
           "S2":"spin2.art.spinNumber", "mu0":"sconst.mu_0", 
           "orientation_x1":"spin1.orientation_axis[0]", "orientation_y1":"spin1.orientation_axis[1]",
           "orientation_z1":"spin1.orientation_axis[2]", 
           "orientation_x2":"spin2.orientation_axis[0]", "orientation_y2":"spin2.orientation_axis[1]",
           "orientation_z2":"spin2.orientation_axis[2]", 
           "relative_x":"spin1.get_relative_coordinates(spin2)[0]",
           "relative_y":"spin1.get_relative_coordinates(spin2)[1]",
           "relative_z":"spin1.get_relative_coordinates(spin2)[2]"}

def load_Hamiltonians(filename):
    f=open(filename, "r")
    text=f.readlines()
    f.close()
    dic={}
    for i in range(1, len(text), 5):
        attributes=text[i].split("\t")
        name=attributes[1]
        dic[name]=text[i-1:i+4]
    return dic
            
H_zee=Hamiltonian()
H_zee.name="H_zee"
#H_zee.equation=[["spin.art.gyro", "B[0]", "Sx"], ["spin.art.gyro", "B[1]", "Sy"], ["spin.art.gyro", "B[2]", "Sz"]]
H_zee.equation=[['gyro', 'B'+i, 'S'+i] for i in ['x', 'y', 'z']]
H_zee.SpinNumber=1

H_zfs=Hamiltonian()
H_zfs.name="H_zfs"
"""
H_zfs.equation=[["D_zfs", "orientation_x^2", "Sx", "Sx"], ["D_zfs", "orientation_x", "orientation_y", "Sx", "Sy"],
                ["D_zfs", "orientation_x", "orientation_z", "Sx", "Sz"], ["D_zfs", "orientation_y", "orientation_x", "Sy", "Sx"],
                ["D_zfs", "orientation_y^2", "Sy", "Sy"], ["D_zfs", "orientation_y", "orientation_z", "Sy", "Sz"],
                ["D_zfs", "orientation_z", "orientation_x", "Sz", "Sx"], ["D_zfs", "orientation_z", "orientation_y", "Sz", "Sy"],
                ["D_zfs", "orientation_z^2", "Sz", "Sz"]]
                """
H_zfs.equation=[['D_zfs', 'orientation_'+i+'^2', 'S'+i, 'S'+i] for i in ['x', 'y', 'z']]
H_zfs.SpinNumber=1

H_dip=Hamiltonian()
H_dip.name="H_dip"
"""
H_dip.equation=[['gyro1', 'gyro2', 'mu0', 'h', '8^(-1)', 'np.pi^(-2)', 'r12^3', 'Sx1', 'Sx2'],
                ['gyro1', 'gyro2', 'mu0', 'h', '8^(-1)', 'np.pi^(-2)', 'r12^3', 'Sy1', 'Sy2'],
                ['gyro1', 'gyro2', 'mu0', 'h', '8^(-1)', 'np.pi^(-2)', 'r12^3', 'Sz1', 'Sz2'],
                ['-3', 'gyro1', 'gyro2', 'mu0', 'h', '8^(-1)', 'np.pi^(-2)', 'r12^5', 'relative_x^2', 'Sx1', 'Sx2'],
                ['-3', 'gyro1', 'gyro2', 'mu0', 'h', '8^(-1)', 'np.pi^(-2)', 'r12^5', 'relative_y^2', 'Sy1', 'Sy2'],
                ['-3', 'gyro1', 'gyro2', 'mu0', 'h', '8^(-1)', 'np.pi^(-2)', 'r12^5', 'relative_z^2', 'Sz1', 'Sz2'],
                ['-3', 'gyro1', 'gyro2', 'mu0', 'h', '8^(-1)', 'np.pi^(-2)', 'r12^5', 'relative_x', 'relative_y', 'Sx1', 'Sy2'],
                ['-3', 'gyro1', 'gyro2', 'mu0', 'h', '8^(-1)', 'np.pi^(-2)', 'r12^5', 'relative_x', 'relative_z', 'Sx1', 'Sz2'],
                ['-3', 'gyro1', 'gyro2', 'mu0', 'h', '8^(-1)', 'np.pi^(-2)', 'r12^5', 'relative_x', 'relative_y', 'Sy1', 'Sx2'],
                ['-3', 'gyro1', 'gyro2', 'mu0', 'h', '8^(-1)', 'np.pi^(-2)', 'r12^5', 'relative_z', 'relative_y', 'Sy1', 'Sz2'],
                ['-3', 'gyro1', 'gyro2', 'mu0', 'h', '8^(-1)', 'np.pi^(-2)', 'r12^5', 'relative_z', 'relative_x', 'Sz1', 'Sx2'],
                ['-3', 'gyro1', 'gyro2', 'mu0', 'h', '8^(-1)', 'np.pi^(-2)', 'r12^5', 'relative_z', 'relative_y', 'Sz1', 'Sy2']]
                """
H_dip.equation=[['gyro1', 'gyro2', 'mu0', 'h', '8^(-1)', 'np.pi^(-2)', 'r12^(-3)', '10^(14)', 'S'+i+'1', 'S'+i+'2'] for i in ['x', 'y', 'z']]+[['-3', 'gyro1', 'gyro2', 'mu0', 'h', '8^(-1)', 'np.pi^(-2)', 'r12^(-5)', 'relative_'+i, 'relative_'+j, 'S'+i+'1', 'S'+j+'2'] for i in ['x', 'y', 'z'] for j in ['x', 'y', 'z']]

H_dip.SpinNumber=2

H_lib = {'zeeman' : H_zee,
         'zfs' : H_zfs,
         'dipolar' : H_dip}
